package com.sxau.apm.monitor.transformer;

import com.sxau.apm.monitor.common.Template;
import javassist.*;

import java.io.IOException;
import java.security.ProtectionDomain;
import java.util.ArrayList;
import java.util.List;

public class ClassFileTransformer implements java.lang.instrument.ClassFileTransformer {

    public static final List<String> INJECTED_CLASS = new ArrayList<>();

    static {
        //被代理的类
        INJECTED_CLASS.add("com.sxau.apm.monitor.AppServlet");
    }

    @Override
    public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined, ProtectionDomain protectionDomain, byte[] classfileBuffer) {
        String realClassName = className.replace("/", ".");

        try {
            if (isToIntercept(realClassName)) {
                ClassPool classPool = ClassPool.getDefault();
                ClassClassPath classPath = new ClassClassPath(this.getClass());
                classPool.insertClassPath(classPath);
                CtClass ctClass = classPool.get(realClassName);
                CtMethod[] declaredMethods = ctClass.getDeclaredMethods();
                for (CtMethod ctMethod : declaredMethods) {
                    String oldMethodName = ctMethod.getName();
                    String threadName = oldMethodName + "_" + ctClass.getName();

                    CtMethod newMethod = CtNewMethod.copy(ctMethod, ctClass, null);
                    ctMethod.setName(oldMethodName + "$agent");
                    String template = ctMethod.getReturnType().getName().equals("void") ? Template.voidFunctionTemplate : Template.returnedFunctionTemplate;
                    newMethod.setBody(String.format("{%s}", String.format(template, threadName, oldMethodName)));
                    ctClass.addMethod(newMethod);
                }

                ctClass.writeFile("/Users/lb/IdeaProjects/apm_monitor/result/");
                return ctClass.toBytecode();
            }
        } catch (CannotCompileException | IOException | NotFoundException e) {
            e.printStackTrace();
            throw new RuntimeException(e);
        }
        return null;
    }

    private static boolean isToIntercept(String realClassName) {
        return INJECTED_CLASS.contains(realClassName);
    }

}
